#!/usr/bin/env python
import numpy as np
from astropy.io import fits
from glob import glob
import sys
from astropy.table import Table
import matplotlib.pyplot as plt
from scipy.stats import sigmaclip

xc = int(sys.argv[2])
yc = int(sys.argv[3])
xs = int(sys.argv[4])
ys = int(sys.argv[5])
l = int(sys.argv[6])

bias_files = glob(sys.argv[1]+"/ptc-0-*.fits")
files = glob(sys.argv[1]+"/ptc-*-*.fits")
dn = []
var = []
for i in range(len(bias_files)):
    if i==0:
        bias = fits.getdata(bias_files[i]).astype('float32')
    else:
        bias += fits.getdata(bias_files[i]).astype('float32')
        
bias /= len(bias_files)

for i in range(l):
    print(i)
    im1 = fits.getdata(sys.argv[1]+"/ptc-"+str(i+1)+"-1.fits").astype('float32')[int(yc-ys/2):int(yc+ys/2),int(xc-xs/2):int(xc+xs/2)]
    im2 = fits.getdata(sys.argv[1]+"/ptc-"+str(i+1)+"-2.fits").astype('float32')[int(yc-ys/2):int(yc+ys/2),int(xc-xs/2):int(xc+xs/2)]
    bias_cut = bias[int(yc-ys/2):int(yc+ys/2),int(xc-xs/2):int(xc+xs/2)]
    dn_raw = (im1+im2)/2-bias_cut
    var_raw = im1-im2
    dn_c,_,_ = sigmaclip(dn_raw.flatten(),low=5,high=5)
    var_c,_,_ = sigmaclip(var_raw.flatten(),low=5,high=5)
    dn.append(np.mean(dn_c))
    var.append(np.var(var_c)/2)

tab = Table()
tab['dn'] = dn
tab['var'] = var
tab.write(sys.argv[1]+"/ptc.tab",format='ipac',overwrite=True)

plt.plot(dn,var,linestyle='--')
plt.scatter(dn,var,s=20)
plt.xlabel("DN")
plt.ylabel("Var")
plt.savefig(sys.argv[1]+"/ptc_result.jpg")
plt.show()
